#include <cstdlib>
#include <cublas_v2.h>
#include <stdio.h>
#include <iostream>

#define CHECK_CUDA(call) {                                   \
    cudaError_t err = call;                                 \
    if (err != cudaSuccess) {                               \
        fprintf(stderr, "CUDA error in file '%s' in line %i: %s.\n", \
                __FILE__, __LINE__, cudaGetErrorString(err)); \
        exit(EXIT_FAILURE);                                 \
    }                                                       \
}


#define CHECK_CUBLAS(call) {                                             \
    cublasStatus_t err = call;                                           \
    if (err != CUBLAS_STATUS_SUCCESS) {                                  \
        fprintf(stderr, "CUBLAS Error in file '%s' in line %i: ",        \
                __FILE__, __LINE__);                                     \
        switch (err) {                                                   \
            case CUBLAS_STATUS_NOT_INITIALIZED:                          \
                fprintf(stderr, "CUBLAS_STATUS_NOT_INITIALIZED\n");      \
                break;                                                   \
            case CUBLAS_STATUS_ALLOC_FAILED:                             \
                fprintf(stderr, "CUBLAS_STATUS_ALLOC_FAILED\n");         \
                break;                                                   \
            case CUBLAS_STATUS_INVALID_VALUE:                            \
                fprintf(stderr, "CUBLAS_STATUS_INVALID_VALUE\n");        \
                break;                                                   \
            case CUBLAS_STATUS_ARCH_MISMATCH:                            \
                fprintf(stderr, "CUBLAS_STATUS_ARCH_MISMATCH\n");        \
                break;                                                   \
            case CUBLAS_STATUS_MAPPING_ERROR:                            \
                fprintf(stderr, "CUBLAS_STATUS_MAPPING_ERROR\n");        \
                break;                                                   \
            case CUBLAS_STATUS_EXECUTION_FAILED:                         \
                fprintf(stderr, "CUBLAS_STATUS_EXECUTION_FAILED\n");     \
                break;                                                   \
            case CUBLAS_STATUS_INTERNAL_ERROR:                           \
                fprintf(stderr, "CUBLAS_STATUS_INTERNAL_ERROR\n");       \
                break;                                                   \
            case CUBLAS_STATUS_NOT_SUPPORTED:                            \
                fprintf(stderr, "CUBLAS_STATUS_NOT_SUPPORTED\n");        \
                break;                                                   \
            case CUBLAS_STATUS_LICENSE_ERROR:                            \
                fprintf(stderr, "CUBLAS_STATUS_LICENSE_ERROR\n");        \
                break;                                                   \
            default:                                                     \
                fprintf(stderr, "UNKNOWN ERROR\n");                      \
        }                                                               \
        exit(EXIT_FAILURE);                                             \
    }                                                                   \
}
int main() {
    cublasHandle_t handle;
    cublasStatus_t stat;
    const int m = 4;
    const int k = 4;
    const int n = 3;
    float *a = new float[m * k];
    for (int i = 0; i < m * k; i++) {
        a[i] = i + 1;
    }
    float b_data[] = {6, 12, 6, 0, 4, 8, 7, 6, 12, 4, 7, 5};
    float *b = new float[k * n];
    float *mem_a = new float[m*k];
    float *mem_b = new float[k*n];
    for (int i = 0; i < k * n; i++) {
        b[i] = b_data[i];
    }
    float c[m * n];
    float *d_a, *d_b, *d_c;
    size_t a_size = m * k * sizeof(float);
    size_t b_size = k * n * sizeof(float);
    size_t c_size = m * n * sizeof(float);
    stat = cublasCreate(&handle);
    if (stat != CUBLAS_STATUS_SUCCESS) {
        const char* e = cublasGetStatusString(stat);
        std::cerr << "初始化CUBLAS失败\n"<<" for "<<e<<"\n";
        return EXIT_FAILURE;
    }
    CHECK_CUDA(cudaMalloc(reinterpret_cast<void**>(&d_a), a_size));
    CHECK_CUDA(cudaMalloc(reinterpret_cast<void**>(&d_b), b_size));
    CHECK_CUDA(cudaMalloc(reinterpret_cast<void**>(&d_c), c_size));

    CHECK_CUBLAS(cublasSetMatrix(m, k, sizeof(float), a, k, d_a, k));
    CHECK_CUDA(cudaMemcpy(d_a,a,a_size,cudaMemcpyHostToDevice));
    // CHECK_CUDA(cudaMemcpy(mem_a,d_a,sizeof(float)*m*k,cudaMemcpyDeviceToHost));
    // for(int i=0;i<m*k;i++){
    //         std::cout<<mem_a[i]<<" ";
    // }
    // CHECK_CUBLAS(cublasSetMatrix(k,n, sizeof(float), b, k, d_b, k));

    // CHECK_CUDA(cudaMemcpy(mem_b,d_b,sizeof(float)*k*n,cudaMemcpyDeviceToHost));
    // std::cout<<"data B = ";
    // for(int i=0;i<n*k;i++){
    //         std::cout<<mem_b[i]<<" ";
    // }
    CHECK_CUDA(cudaMemcpy(d_b,b,b_size,cudaMemcpyHostToDevice));
    if (stat != CUBLAS_STATUS_SUCCESS) {
        std::cerr << "初始化GEMM内存失败\n";
        cudaFree(d_a);
        cublasDestroy(handle);
        return EXIT_FAILURE;
    }
    const float alpha = 1.f;
    const float beta = 0.f;
    CHECK_CUBLAS(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, d_a, m, d_b, k, &beta, d_c, m));//列优先存储
    CHECK_CUBLAS(cublasGetMatrix(m, n, sizeof(float), d_c, m, c, m));
    std::cout<<"\n";
    for (int i = 0; i < m * n; i++) {
        std::cout << "data[" << i << "] = " << c[i] << "\n";
    }
    cudaFree(d_a);
    cudaFree(d_b);
    cudaFree(d_c);
    cublasDestroy_v2(handle);

}
